#!/usr/bin/perl

use strict;
use LWP;
use Digest::MD5;
use Getopt::Long;
use DBI;

use constant VERSION => "0.1b";
use constant SCHEMA => 1;

use constant ERR_USAGE => 100;
use constant ERR_CHECKER => 101;

use constant VIRUS_FOUND => 1;
use constant VIRUS_NOT_FOUND => 0;

my $VERSION = VERSION;
my $AUTHOR = "Iñaki Rodríguez - http://www.virtualminds.es (2009)";

# Global Variables
my $md5 = Digest::MD5->new;
my $virhash;
my $file;
my $hash;
my $v = 0;
my $quiet = 0;
my $nocache = 0;
my $initial = 0;
my $result;
my $source;
my $ecode = VIRUS_NOT_FOUND;

# See README.sql for table definition
# or use --initial 
# SQlite # REMEMBER: Change db location.
my %attr = (
    PrintError => 0,
    RaiseError => 0,
);

my $dbname = 'virus.db';
my $dbuser = "";
my $dbpass = "";
my $dbh;
if(!$nocache) {
	$dbh = DBI->connect("dbi:SQLite:dbname=$dbname",$dbuser,$dbpass,\%attr);
	die ("Cannot open SQL connection") if !$dbh;
}

# Or better, store it in separate config file.
if( -f "./malware_hash.conf") { require "./malware_hash.conf"; } 


my %hash_checkers = ( 	"virustotal" => \&check_virustotal,
						"dummy" => \&check_dummy
					);
				

# Options
GetOptions ( 
				"file=s" => \$file,
				"hash=s" => \$hash,
				"verbose|v" => \$v,
				"no-cache" => \$nocache,
				"initial" => \$initial,
				"quiet" => \$quiet
);

if($initial) { 
	db_initial_schema();
	exit 0;
}

if ($v && $quiet) {
	usage();
	exit ERR_USAGE;
}

if($file && $hash) {
	usage();
	exit ERR_USAGE;
}

if(!$file && !$hash) {
	usage();
	exit ERR_USAGE;
}


if(-f $file) {
	open FILE,$file;
	binmode(FILE);
	$virhash = $md5->addfile(*FILE)->hexdigest;
	print "[HASH] Using MD5 hash $virhash\n" if $v;
}
elsif ($hash && length($hash) == 32) {
	$virhash = $hash;
	print "[HASH] Using hash $virhash\n" if $v;
} else {
	usage();
	exit ERR_USAGE;
}

if (!$nocache) { # Search previously result in blacklist
	if(db_get_blacklist($virhash)) { 
		print "[CACHED] Virus Found\n";
		exit VIRUS_FOUND;
	}
}

foreach $source (keys %hash_checkers) {

	my $cmd = $hash_checkers{$source};
	$result = \&$cmd($virhash);
	if($$result == VIRUS_FOUND) { 
		print "[$source] Virus Found\n";
		db_insert_blacklist($virhash,$source) if (!$nocache);
		$ecode = VIRUS_FOUND;
	} else {
		print "[$source] Virus not found\n" if($v);
	}
}

exit $ecode;


### CHECKER ###

sub check_virustotal {

	my $ua = LWP::UserAgent->new;
	push @{$ua->requests_redirectable }, 'POST';

	my $resp = $ua->post('http://www.virustotal.com/vt/en/consultamd5',[ "hash" => $virhash , "x" => 138 , "y" => 24 ]);

	if($resp->is_success) {
		my $data = $resp->content;
		$data =~ /status_porcentaje(.*)/;
		my $virdata = $1;
		$virdata =~ m,>(\d+)</span>/(\d+) ,;
		if($1 > 0) {
			return VIRUS_FOUND;
		} else {
			return VIRUS_NOT_FOUND;
		}
	} else {
		return ERR_CHECKER;
	}
}

sub usage () {
	print << "_EOT_"

	Malware Hash Checker v$VERSION - $AUTHOR
	--------------------------------
	$0 --file filename | --hash md5 hash [-v|-q|--no-cache]
	
	$0 --initial

Arguments:

	-q			Quiet mode
	-v			Verbose mode
	--no-cache		Don't use cache
	--initial		Create initial database (for cache support)

_EOT_
}

sub db_insert_blacklist {
	my $hash = shift;
	my $source = shift;

	my $res = $dbh->prepare("INSERT INTO blacklist(md5hash,source,created) values (?,?,?)");
	$res->execute($hash,$source,time());
	$res->finish();
}

sub db_get_blacklist {
	my $hash = shift;

	my $res = $dbh->prepare("SELECT md5hash from blacklist where md5hash = ?");
	$res->execute($hash);
    if($res->fetchrow) {
		$res->finish();
        return VIRUS_FOUND;
    } else {
		$res->finish();
        return VIRUS_NOT_FOUND;
    }
}

sub db_initial_schema {
	$dbh->do("BEGIN TRANSACTION") || die ("ERROR: Cannot create database correctly!");
	$dbh->do("CREATE TABLE conf (key varchar(50) not null primary key, value varchar(100) not null)")  || roll_and_die ("ERROR: Cannot create database correctly!");
	$dbh->do("CREATE TABLE blacklist (md5hash varchar(32) not null primary key, source varchar(80), created timestamp)") || roll_and_die ("ERROR: Cannot create database correctly!"); 
	my $res = $dbh->prepare("INSERT INTO conf(key,value) values (\"schema_version\",?)");
	$res->execute(SCHEMA) || roll_and_die ("ERROR: Cannot create database correctly!"); 
	$dbh->do("COMMIT");
}

sub roll_and_die {
	my $msg = shift;

	$dbh->do("ROLLBACK");
	die($msg);
}

sub check_dummy {
	return VIRUS_NOT_FOUND;
}
